"""
python p2_solution_pygad_genetic.py
"""

import json
import os
import pygad
import numpy as np

# 路径设置
warehouse_path = '../fujian/fujian3/origin_data/warehouse.json'
inventory_path = '../fujian/fujian3/data_from_p1/all_average_inventory.json'
sales_path = '../fujian/fujian3/data_from_p1/all_average_sales.json'
output_path = '../fujian/fujian3/pygad/'

# 数据加载
with open(warehouse_path, 'r') as f:
    warehouses = json.load(f)

with open(inventory_path, 'r') as f:
    inventory_data = json.load(f)

with open(sales_path, 'r') as f:
    sales_data = json.load(f)

# 数据处理
warehouse_dict = {w['warehouse_id']: w for w in warehouses}
inventory_dict = {i['category_id']: i['average_inventory'] for i in inventory_data}
sales_dict = {s['category_id']: s['average_sales'] for s in sales_data}

categories = list(inventory_dict.keys())
num_categories = len(categories)
warehouses_list = list(warehouse_dict.keys())

# 计算成本和约束
def calculate_cost(solution):
    total_cost = 0
    warehouse_utilization = {w: {'inventory': 0, 'sales': 0} for w in warehouses_list}
    
    for category_index, warehouse_index in enumerate(solution):
        category_id = categories[category_index]
        warehouse_id = warehouses_list[warehouse_index]

        inventory = inventory_dict[category_id]
        sales = sales_dict[category_id]

        warehouse_utilization[warehouse_id]['inventory'] += inventory
        warehouse_utilization[warehouse_id]['sales'] += sales

        if inventory > 0:
            total_cost += warehouse_dict[warehouse_id]['daily_cost']
    
    # 检查约束条件
    for warehouse_id, usage in warehouse_utilization.items():
        if (usage['inventory'] >= warehouse_dict[warehouse_id]['max_inventory'] or
            usage['sales'] >= warehouse_dict[warehouse_id]['max_sales']):
            return float('inf'), {}  # 违反约束条件，返回无穷大
    
    return total_cost, warehouse_utilization

# 修正适应度函数
def fitness_func(ga_instance, solution, solution_idx):
    cost, _ = calculate_cost(solution)
    return 1.0 / (cost + 1)  # 适应度函数

# 初始化遗传算法
num_generations = 100
num_parents_mating = 10

sol_per_pop = 20
initial_population = np.random.randint(0, len(warehouses_list), size=(sol_per_pop, num_categories))

# 创建遗传算法对象
ga_instance = pygad.GA(num_generations=num_generations,
                        num_parents_mating=num_parents_mating,
                        fitness_func=fitness_func,
                        sol_per_pop=sol_per_pop,
                        initial_population=initial_population,
                        num_genes=num_categories,
                        gene_type=int,
                        mutation_probability=0.1)

# 运行遗传算法
ga_instance.run()

# 获取最佳解
best_solution, best_solution_fitness, _ = ga_instance.best_solution()
total_cost, warehouse_utilization = calculate_cost(best_solution)

# 输出结果
assignment = [{'category_id': categories[i], 'warehouse_id': warehouses_list[best_solution[i]]} for i in range(num_categories)]
os.makedirs(output_path, exist_ok=True)

# 保存分配方案
with open(os.path.join(output_path, 'assignment.json'), 'w') as f:
    json.dump(assignment, f)

# 保存总成本
with open(os.path.join(output_path, 'all_cost.txt'), 'w') as f:
    f.write(str(total_cost))

# 计算和保存利用率
utilization_rates_inventory = []
utilization_rates_sales = []

for warehouse_id, usage in warehouse_utilization.items():
    utilization_rates_inventory.append({
        'warehouse_id': warehouse_id,
        'utilization_rate_of_inventory': usage['inventory'] / warehouse_dict[warehouse_id]['max_inventory']
    })
    utilization_rates_sales.append({
        'warehouse_id': warehouse_id,
        'utilization_rate_of_sales': usage['sales'] / warehouse_dict[warehouse_id]['max_sales']
    })

# 保存库存利用率
with open(os.path.join(output_path, 'all_utilization_rate_inventory.json'), 'w') as f:
    json.dump(utilization_rates_inventory, f)

# 保存出货量利用率
with open(os.path.join(output_path, 'all_utilization_rate_sales.json'), 'w') as f:
    json.dump(utilization_rates_sales, f)

print("完成！最佳总成本:", total_cost)
